# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import os
import time
import uuid
import json
from typing import List
from fastapi import UploadFile
from tortoise.expressions import Q

from common.enums.agent import AgentEnum
from exception import AppException
from config import get_settings
from common.utils.urls import UrlUtil
from common.utils.times import TimeUtil
from common.utils.tools import ToolsUtil
from common.utils.files import FilesUtil
from common.utils.config import ConfigUtil
from common.enums.models import ModelsEnum
from common.models.dev import DevModelsModel
from common.models.docs import DocsFilesModel
from common.models.docs import DocsPromptModel
from common.models.docs import DocsRecordsModel
from common.postgres.public import PgDocumentsModel
from common.cache.pool_cache import KeyPoolCache
from plugins.storage.driver import StorageDriver
from apps.api.schemas.apply.docs import files_schema as schema


class DocsFilesService:

    @classmethod
    async def prompts(cls) -> List[schema.FilesPromptsVo]:
        """
        文件模式

        Returns:
            List[schema.FilesPromptsVo]: 模式列表Vo

        Author:
            zero
        """
        lists = await (DocsPromptModel
                       .filter(is_delete=0)
                       .filter(is_disable=0)
                       .all()
                       .order_by("-sort", "-id")
                       .values("code", "icon", "name", "intro"))

        _lists = []
        for item in lists:
            _lists.append(schema.FilesPromptsVo(
                code=item["code"],
                icon=await UrlUtil.to_absolute_url(item["icon"]),
                name=item["name"],
                intro=item["intro"],
            ))
        return _lists

    @classmethod
    async def lists(cls, params: schema.FilesSearchIn, user_id: int) -> List[schema.FilesListVo]:
        """
        文件列表

        Args:
            params (schema.FilesSearchIn): 搜索参数
            user_id (int): 用户ID

        Returns:
            List[schema.FilesListVo]: 文件列表Vo

        Author:
            zero
        """
        # 查询条件
        where = [Q(user_id=user_id), Q(is_delete=0)]
        if params.keyword:
            where.append(Q(name__icontains=params.keyword))

        # 查询数据
        lists = await DocsFilesModel.filter(*where).order_by("-id").limit(200)

        # 处理数据
        _lists = []
        for item in lists:
            _lists.append(schema.FilesListVo(
                code=item.code,
                name=item.name,
                size=item.size,
                path=await UrlUtil.to_absolute_url(item.pdf_path),
                create_time=TimeUtil.timestamp_to_date(item.create_time)
            ))

        return _lists

    @classmethod
    async def detail(cls, code: str, user_id: int) -> schema.FilesDetailVo:
        """
        文件详情

        Args:
            code (str): 文件编号
            user_id (int): 用户ID

        Returns:
            List[schema.FilesDetailVo]: 文件详情Vo

        Author:
            zero
        """
        # 查询文件
        files = await DocsFilesModel.filter(code=code, user_id=user_id).first()
        if files is None:
            raise AppException("找不到文件")
        if files.is_delete:
            raise AppException("文件已被删除")

        # 查询模式
        prompt_icon = ""
        prompt_name = ""
        prompt_intro = ""
        prompt = await DocsPromptModel.filter(code=files.mode).first()
        if prompt is not None:
            prompt_icon = await UrlUtil.to_absolute_url(prompt.icon)
            prompt_name = prompt.name
            prompt_intro = prompt.intro

        # 查询配置
        auto_abstract = await ConfigUtil.get("docs", "auto_abstract", "0")
        chat_avatar = await ConfigUtil.get("docs", "chat_avatar", "")
        chat_avatar = await UrlUtil.to_absolute_url(chat_avatar)

        # 查询记录
        results = await (DocsRecordsModel
                         .filter(scene="abstract")
                         .filter(file_id=files.id)
                         .filter(user_id=user_id)
                         .filter(is_clear=0)
                         .filter(is_delete=0)
                         .order_by("-id")
                         .first().values("answer"))

        abstract = (results.get("answer") or "") if results else ""

        # 返回结果
        return schema.FilesDetailVo(
            code=code,
            name=files.name,
            size=files.size,
            path=await UrlUtil.to_absolute_url(files.pdf_path),
            error=files.error or "",
            status=files.status,
            auto_abstract=int(auto_abstract),
            chat_avatar=chat_avatar,
            prompt_icon=prompt_icon,
            prompt_name=prompt_name,
            prompt_intro=prompt_intro,
            create_time=TimeUtil.timestamp_to_date(files.create_time),
            abstract=abstract
        )

    @classmethod
    async def rename(cls, post: schema.FilesRenameIn, user_id: int):
        """
        文件命名

        Args:
            post (schema.FilesRenameIn): 更新参数
            user_id (int): 用户ID

        Author:
            zero
        """
        files = await (DocsFilesModel
                       .filter(code=post.code)
                       .filter(user_id=user_id)
                       .first())

        if not files:
            raise AppException("文件不存在")
        if files.is_delete:
            raise AppException("文件已被删除")

        await DocsFilesModel.filter(id=files.id).update(
            name=post.name,
            update_time=int(time.time())
        )

    @classmethod
    async def delete(cls, code: str, user_id: int):
        """
        文件删除

        Args:
            code (str): 文档编号
            user_id (int): 用户ID

        Author:
            zero
        """
        files = await (DocsFilesModel
                       .filter(code=code)
                       .filter(user_id=user_id)
                       .filter(is_delete=0)
                       .first())

        if not files:
            raise AppException("文件不存在")
        if files.is_delete:
            raise AppException("文件已被删除")

        await DocsFilesModel.filter(id=files.id).update(
            is_delete=1,
            delete_time=int(time.time())
        )

        await DocsRecordsModel.filter(file_id=files.id).update(
            is_delete=1,
            delete_time=int(time.time())
        )

        await StorageDriver.delete(files.raw_path)
        await StorageDriver.delete(files.pdf_path)
        await StorageDriver.delete(files.txt_path)

    @classmethod
    async def retry(cls, code: str, user_id: int):
        """
        失败重试

        Args:
            code (str): 文档编号
            user_id (int): 用户ID

        Author:
            zero
        """
        files = await (DocsFilesModel
                       .filter(code=code)
                       .filter(user_id=user_id)
                       .filter(is_delete=0)
                       .first())

        if not files:
            raise AppException("文件不存在")
        if files.is_delete:
            raise AppException("文件已被删除")
        if files.status != AgentEnum.EMB_FAIL:
            raise AppException("不是失败状态不可操作")

        await DocsFilesModel.filter(id=files.id).update(
            error="",
            status=AgentEnum.EMB_WAIT,
            update_time=int(time.time())
        )

    @classmethod
    async def upload(cls, file_in: UploadFile, mode: str, user_id: int, ip: str) -> schema.FilesUploadVo:
        """
        文件上传

        Args:
            file_in (UploadFile): 上传参数
            mode (str): 上传模式
            user_id (int): 用户ID
            ip (int): 操作IP

        Author:
            zero
        """
        raw_file_path = None
        pdf_file_path = None
        txt_file_path = None
        try:
            # 获取存在位置
            UP_ROOT: str = str(get_settings().UPLOAD.get("root"))
            UP_STOR: str = str(get_settings().UPLOAD.get("path"))

            # 验证文件模式
            prompt = await DocsPromptModel.filter(code=mode, is_delete=0).first()
            if prompt is None:
                raise AppException("选择的模式不存在")
            if prompt.is_disable:
                raise AppException("选择的模式被禁用")

            # 验证向量模型
            vector_models = await ConfigUtil.get("docs", "vector_models") or ""
            vector = await DevModelsModel.check_models(models=vector_models, scene=ModelsEnum.TYPE_VECTOR)

            # 验证密钥配置
            key = await KeyPoolCache(scene=ModelsEnum.TYPE_VECTOR, channel=vector["mark"]).get_key()
            if not key:
                raise AppException("当前应用尚未配置相关密钥")

            # 文件上传保存
            result = await StorageDriver.upload(
                file_in,
                root=UP_ROOT,
                stores=UP_STOR,
                folder="applies/docs/raw",
                synchro=True
            )

            # 提取文件路径
            rel_pdf_path = result.get("path")
            raw_file_path = UrlUtil.to_root_path(result.get("path"))
            pdf_file_path = raw_file_path

            # 转换pdf文件
            if result["ext"] != "pdf":
                pdf_file_name = result["path"].replace("docs/raw", "docs/pdf")
                pdf_file_name = pdf_file_name.replace(result["ext"], "pdf")
                pdf_file_path = UrlUtil.to_root_path(pdf_file_name)
                rel_pdf_path = pdf_file_name
                # 发起Soffice转换
                await FilesUtil.convert_file(raw_file_path, output_path=os.path.dirname(pdf_file_path))
                # 本地文件上传至OSS
                await StorageDriver.push(path=pdf_file_path, key=rel_pdf_path)

            # 提取文件内容
            pages, chunks = await FilesUtil.parse_pdf(pdf_file_path, chunk_size=512)

            # 记录到数据库
            salt: str = str(time.time()) + str(user_id) + str(result["path"])
            code: str = ToolsUtil.make_md5_str(salt)
            file = await DocsFilesModel.create(
                code=code,
                mode=mode,
                user_id=user_id,
                name=result["name"],
                size=result["size"],
                ext=result["ext"],
                raw_path=result["path"],
                pdf_path=rel_pdf_path,
                txt_path="",
                vector_models=vector_models,
                vector_alias=vector.get("alias"),
                ip=ip,
                create_time=int(time.time()),
                update_time=int(time.time())
            )

            # 异步写入文本
            folder_date = TimeUtil.timestamp_to_date(int(time.time()), formats="%Y%m%d")
            txt_file_name = f"perusal/{folder_date}/{file.id}.md"
            txt_file_path = UrlUtil.to_runtime_path(txt_file_name)
            json_text_str = json.dumps(pages, indent=4, ensure_ascii=False)
            await ToolsUtil.write_file(txt_file_path, json_text_str)

            # 更新解析路径
            await DocsFilesModel.filter(id=file.id).update(txt_path=txt_file_name)

            # 批量处理文本
            pg_uuid = []
            pg_docs = []
            for chunk in chunks:
                _uuid = uuid.uuid4()
                pg_uuid.append(str(_uuid))
                pg_docs.append(PgDocumentsModel(
                    uuid=_uuid,
                    user_id=user_id,
                    file_id=file.id,
                    vector_model=vector_models,
                    vector_alias=vector.get("alias"),
                    page_no=chunk["page_no"],
                    page_nv=",".join(chunk["page_nv"]),
                    chunk_index=chunk["chunk_index"],
                    question=chunk["chunk_texts"],
                    create_time=int(time.time()),
                    update_time=int(time.time())
                ))

            # 写入到PG库
            if pg_docs:
                await PgDocumentsModel.bulk_create(pg_docs)
                await PgDocumentsModel.update_ts_vectory(pg_uuid)

            # 返回上传结果
            return schema.FilesUploadVo(
                code=code,
                name=result["name"],
                size=result["size"],
                ext=result["ext"],
                path=await UrlUtil.to_absolute_url(result["path"])
            )
        except Exception as e:
            if raw_file_path and os.path.exists(raw_file_path):
                os.unlink(raw_file_path)
            if pdf_file_path and os.path.exists(pdf_file_path):
                os.unlink(pdf_file_path)
            if txt_file_path and os.path.exists(txt_file_path):
                os.unlink(txt_file_path)
            raise AppException(str(e))
        finally:
            engine: str = await UrlUtil.get_storage_engine()
            if engine != "local":
                if raw_file_path and os.path.exists(raw_file_path):
                    os.unlink(raw_file_path)
                if pdf_file_path and os.path.exists(pdf_file_path):
                    os.unlink(pdf_file_path)
